import numpy as np
import pandas as pd
import copy
import fun
import itertools

# 定义常量
ZoneMax = 54
CropMax = 41
Is_Bean_Cons = False

# 参数调整
loss_degree = 0  # 亏损程度enum(0, 0.5)
sale_p = 0.5  # 小麦、玉米，预期销售量增长率 5% ~ 10%
ge_sele_p = 0  # 其他作物预期销量增长率 -5% ~ 5%
yield_P = 0  # 亩产量变化 -10% ~ 10%
cost_p = 0.5  # 种植成本 5%左右
price_p = 0  #

# 读入数据
X_t = pd.read_csv('数据/Xij-2023.csv', index_col=0)  # 表示i号地中j号作物在第t年的种植面积
P = pd.read_csv('数据/Pij.csv', index_col=0)  # 表示i号地中是否可以种植j号作物
C_t = pd.read_csv('数据/Cij.csv', index_col=0)  # 表示i号地中j号作物在第t年的种植成本
Q_t = pd.read_csv('数据/Qij.csv', index_col=0)  # 表示i号地中j号作物在第t年的产量
Z = pd.read_csv('数据/Zi.csv', index_col=0).iloc[0]  # 表示i号地的面积
Price_t = pd.read_csv('数据/Priceij.csv', index_col=0)  # 表示i号地中j号作物在第t年的单价
Need_t = pd.read_csv('数据/Needj.csv', index_col=0).iloc[0]  # 表示j号作物在t年的需求量
X_Pre = pd.DataFrame(copy.copy(X_t))
X_PrePre = pd.DataFrame(copy.copy(X_t))

is_single = pd.read_csv('数据/is_single.csv', index_col=0)


# 返回这块地占用的总面积（考虑单双季）
def occ_area(z):
    s = 0
    for i, x in enumerate(z):
        s += x
        if not is_single.loc[i, '是否单季']:
            s -= x / 2
    return s


# 总面积约束
def area_cons():
    a = [fun.occ_area(X_t.iloc[i]) for i in range(ZoneMax)]
    for idx, s, z in zip(range(len(a)), a, Z):
        if s > z:
            return False
    return True


# 重茬种植约束
def re_cons():
    for i in range(ZoneMax):
        for j in range(CropMax):
            up = P.iloc[i, j] * (Z.iloc[i] - X_Pre.iloc[i, j])
            if not (0 <= X_t.iloc[i, j] <= up):
                return False
    return True


# 三年种豆约束
def bean_cons():
    if not Is_Bean_Cons:
        return True
    for i in range(ZoneMax):
        down = Z.iloc[i] + sum(
            [X_Pre.iloc[i, j] + X_PrePre.iloc[i, j] for j in itertools.chain(range(0, 5), range(16, 19))])
        val = sum([X_t.iloc[i, j] for j in itertools.chain(range(0, 5), range(16, 19))])
        if not (down <= val <= Z.iloc[i]):
            return False
    return True


# 水浇地约束1
def water_area_cons1():
    for i in range(26, 34):
        val = X_t.iloc[i, 15] + sum([X_t.iloc[i, j] for j in range(16, 37)])
        up = 2 * Z.iloc[i]
        if not (0 <= val <= up):
            return False
    return True


# 水浇地约束2
def water_area_cons2():
    for i in range(26, 34):
        val = sum([X_t.iloc[i, j] for j in range(16, 34)])
        if not (0 <= val <= Z.iloc[i]):
            return False
    return True


# 水浇地约束3
def water_area_cons3():
    for i in range(26, 34):
        val = sum([X_t.iloc[i, j] for j in range(34, 37)])
        if not (0 <= val <= Z.iloc[i]):
            return False
    return True


def cons_total():
    if area_cons() and re_cons() and bean_cons() and water_area_cons1() and water_area_cons2() and water_area_cons3():
        return True
    return False


# 目标函数 - 滞销浪费
def Y1():
    global loss_degree
    loss_degree = 0
    a = (X_t * Q_t * Price_t - X_t * C_t).sum().sum()
    b = 0
    for j in range(CropMax):
        for i in range(ZoneMax):
            b += max(X_t.iloc[i, j] * Q_t.iloc[i, j] - Need_t.iloc[j], 0)
    return a - b * (1 - loss_degree)


# 目标函数 - 50%降价
def Y2():
    global loss_degree
    loss_degree = 0.5
    a = (X_t * Q_t * Price_t - X_t * C_t).sum().sum()
    b = 0
    for j in range(CropMax):
        for i in range(ZoneMax):
            b += max(X_t.iloc[i, j] * Q_t.iloc[i, j] - Need_t.iloc[j], 0)
    return a - b * (1 - loss_degree)


def model1():
    global X_t, Is_Bean_Cons, X_Pre, X_PrePre
    gen = fun.generate_matrix_combinations(Z.values, ZoneMax, CropMax)
    for i in range(2024 , 2030):
        mx = 0.0
        if i >= 2025:
            Is_Bean_Cons = True
        for matrix in gen:
            temp = X_t
            X_t = pd.DataFrame(matrix)
            if not cons_total():
                X_t = temp
                continue
            mx = max(mx, Y1())
        X_t.to_csv('./' + str(i) + '.csv')
        X_PrePre = X_Pre
        X_Pre = X_t


def model2():
    global X_t, Is_Bean_Cons, X_Pre, X_PrePre
    gen = fun.generate_matrix_combinations(Z.values, ZoneMax, CropMax)
    for i in range(2024 , 2030):
        mx = 0.0
        if i >= 2025:
            Is_Bean_Cons = True
        for matrix in gen:
            temp = X_t
            X_t = pd.DataFrame(matrix)
            if not cons_total():
                X_t = temp
                continue
            mx = max(mx, Y1())
        X_t.to_csv('./' + str(i) + '.csv')
        X_PrePre = X_Pre
        X_Pre = X_t
